from updateValue import updateValue
import matplotlib.pyplot as plt
import matplotlib.patches as patches  
import numpy as np

class ValueIteration(updateValue):
    def __init__(self, V, env, policy_matrix, action_dict, delta, gamma):
        super().__init__(policy_matrix, env, action_dict, gamma)
        self.V = V
        self.delta = delta
        self.v_texts = {}
        self.canvas = None
        self.ax = None

    def run(self):  
        while True:
        # for x in range(3):
            V_old = self.V.copy()
            self.V = self.update_State_value(self.V)
            self.update_policy_matrix(self.V)
        
            self.render()
            if max([(abs(self.V[V_value] - V_old[V_value])) for V_value in range(len(self.V))]) <= self.delta:
                return self.V, self.policy_matrix
            

    def render(self):
        if self.canvas is None:
            plt.ion()                             
            self.canvas, self.ax = plt.subplots()   
            self.ax.set_xlim(-0.5, self.env.env_size[0] - 0.5)
            self.ax.set_ylim(-0.5, self.env.env_size[1] - 0.5)
            self.ax.xaxis.set_ticks(np.arange(-0.5, self.env.env_size[0], 1))     
            self.ax.yaxis.set_ticks(np.arange(-0.5, self.env.env_size[1], 1))     
            self.ax.grid(True, linestyle="-", color="gray", linewidth="1", axis='both')          
            self.ax.set_aspect('equal')
            self.ax.invert_yaxis()                           
            self.ax.xaxis.set_ticks_position('top')           
            
            idx_labels_x = [i for i in range(self.env.env_size[0])]
            idx_labels_y = [i for i in range(self.env.env_size[1])]
            for lb in idx_labels_x:
                self.ax.text(lb, -0.75, str(lb+1), size=10, ha='center', va='center', color='black')           
            for lb in idx_labels_y:
                self.ax.text(-0.75, lb, str(lb+1), size=10, ha='center', va='center', color='black')
            self.ax.tick_params(bottom=False, left=False, right=False, top=False, labelbottom=False, labelleft=False,labeltop=False)   

            target_rect = patches.Rectangle( (self.env.target_state[0]-0.5, self.env.target_state[1]-0.5), 1, 1, linewidth=1, edgecolor=self.env.color_target, facecolor=self.env.color_target)
            self.ax.add_patch(target_rect)     

            for forbidden_state in self.env.forbidden_states:
                rect = patches.Rectangle((forbidden_state[0]-0.5, forbidden_state[1]-0.5), 1, 1, linewidth=1, edgecolor=self.env.color_forbid, facecolor=self.env.color_forbid)
                self.ax.add_patch(rect)

            for x in range(self.env.env_size[0]):
                for y in range(self.env.env_size[1]):
                    state_id = y * self.env.env_size[1] + x
                    v_text = self.ax.text(x, y, "", size=9, ha='center', va='center', color='black')
                    self.v_texts[state_id] = v_text

        # self.agent_circle.center = (self.agent_state[0], self.agent_state[1])
        for x in range(self.env.env_size[0]):
            for y in range(self.env.env_size[1]):
                state_id = y * self.env.env_size[1] + x
                if state_id in self.v_texts:
                    v_value = self.V[state_id]  # 假设 self.V 是一维数组，按位置存
                    self.v_texts[state_id].set_text(f"{v_value:.2f}")

        plt.draw()
        plt.pause(0.2)

